from torch import nn

class LSTM_v3(nn.Module):
    def __init__(self, dropout=0.25):
        super(LSTM_v3, self).__init__()
        INT_HIDDEN_SIZE = 128  # 进一步增加隐藏层大小
        INT_NUM_LAYERS = 3  # 增加LSTM层数
        CONST_INT_INPUT_SIZE = 5  # [x, y, theta, lstim, rstim]
        CONST_INT_OUTPUT_SIZE = 3  # [x, y, theta]
        CONST_BOOL_BATCH_FIRST = True
        
        # 添加一个双向LSTM来捕获序列中的双向信息
        self.bilstm = nn.LSTM(
            input_size=CONST_INT_INPUT_SIZE,
            hidden_size=INT_HIDDEN_SIZE // 2,
            num_layers=1,
            batch_first=CONST_BOOL_BATCH_FIRST,
            dropout=dropout if INT_NUM_LAYERS > 1 else 0,
            bidirectional=True
        )
        
        self.lstm = nn.LSTM(
            input_size=INT_HIDDEN_SIZE,
            hidden_size=INT_HIDDEN_SIZE,
            num_layers=INT_NUM_LAYERS,
            batch_first=CONST_BOOL_BATCH_FIRST,
            dropout=dropout if INT_NUM_LAYERS > 1 else 0  # 多层LSTM时添加dropout
        )
        
        # 增加多层全连接网络
        self.fc1 = nn.Linear(INT_HIDDEN_SIZE, INT_HIDDEN_SIZE)
        self.ln1 = nn.LayerNorm(INT_HIDDEN_SIZE)  # 使用层归一化代替批归一化
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(dropout)
        
        self.fc2 = nn.Linear(INT_HIDDEN_SIZE, INT_HIDDEN_SIZE // 2)
        self.ln2 = nn.LayerNorm(INT_HIDDEN_SIZE // 2)  # 使用层归一化代替批归一化
        self.relu2 = nn.ReLU()
        self.dropout2 = nn.Dropout(dropout)
        
        self.fc3 = nn.Linear(INT_HIDDEN_SIZE // 2, CONST_INT_OUTPUT_SIZE)

    def forward(self, x):
        # 先经过双向LSTM提取基础特征
        bilstm_out, _ = self.bilstm(x)
        
        # 再经过多层LSTM深入提取序列特征
        lstm_out, _ = self.lstm(bilstm_out)
        last_output = lstm_out[:, -1, :]
        
        # 多层全连接网络处理
        x = self.fc1(last_output)
        x = self.ln1(x)  # 层归一化不受批量大小限制
        x = self.relu(x)
        x = self.dropout(x)
        
        x = self.fc2(x)
        x = self.ln2(x)  # 层归一化不受批量大小限制
        x = self.relu2(x)
        x = self.dropout2(x)
        
        output = self.fc3(x)
        return output

def get_model():
    return LSTM_v3()